import subprocess as sp

# little_cpu_freq_list = [300000, 576000, 748800, 998400, 1209600, 1324800, 1516800, 1612800, 1708800]
# big_cpu_freq_list = [300000, 652800, 825600, 979200, 1132800, 1363200, 1536000, 1747200, 1843200, 1996800]

# TODO: 把频率最低的档位给去掉，调不了那么低
class CPU:
	# The number of CPUs on the device
	n_cpu = 8
	# The temperature info collected during runtime
	temp_data = []
	# The index range of CPUs with different frequencies
	# cpu_idx = {'l': [0, 4], 'm': [4, 7], 'b': [7, 8]}
	
	########################################
	# CPU instances initialization methods #
	########################################
	def __init__(self, idx: int, ip: str):
		self.idx 	   = idx
		self.ip 	   = ip
		self.freq_data = []

		self.freq_list    = self.init_available_freq()
		self.max_freq_idx = len(self.freq_list) - 1
		self.freq_idx     = self.max_freq_idx

		self.init_scaling_range() # TODO: not always minimum

	def init_available_freq(self):
		'''
		Initialize the available frequency list of current CPU instance
		'''
		fname  = f'/sys/devices/system/cpu/cpu{self.idx}/cpufreq/scaling_available_frequencies'
		output = sp.check_output(['adb', '-s', self.ip, 'shell', 'su -c', '\"cat', fname+'\"'])
		# print(f'[c{self.idx}] available frequencies: {output.decode("utf-8").strip()}')
		return [int(f) for f in output.split()]
	
	def init_scaling_range(self):
		'''
		Initialize the frequency scaling range of current CPU instance
		'''
		# Set scaling min frequency
		fname = f'/sys/devices/system/cpu/cpu{self.idx}/cpufreq/scaling_min_freq'
		sp.check_output(['adb', '-s', self.ip, 'shell', 'su -c', '\"echo', str(self.freq_list[0])+' >', fname+'\"'])
		output = sp.check_output(['adb', '-s', self.ip, 'shell', 'su -c', '\"cat', fname+'\"'])
		min_freq = int(output.decode('utf-8').strip())
		# Set scaling max frequency
		fname = f'/sys/devices/system/cpu/cpu{self.idx}/cpufreq/scaling_max_freq'
		sp.check_output(['adb', '-s', self.ip, 'shell', 'su -c', '\"echo', str(self.freq_list[self.max_freq_idx])+' >', fname+'\"'])
		output = sp.check_output(['adb', '-s', self.ip, 'shell', 'su -c', '\"cat', fname+'\"'])
		max_freq = int(output.decode('utf-8').strip())
		# Print prompt
		# print(f'[c{self.idx}] scaling from {min_freq} to {max_freq}')
	
	#####################################
	# CPU instances information methods #
	#####################################
	def set_freq(self, action_idx: int, action_max):
		'''
		Set the frequency of current CPU instance
		'''
		# Compute the index of frequency level
		self.freq_idx = int(action_idx / action_max * (len(self.freq_list) - 1))
		fname = f'/sys/devices/system/cpu/cpu{self.idx}/cpufreq/scaling_setspeed'
		sp.check_output(['adb', '-s', self.ip, 'shell', 'su -c', '\"echo', str(self.freq_list[self.freq_idx])+' >', fname+'\"'])
		
	def log_freq(self):
		'''
		Log the frequency of current CPU instance
		'''
		fname  = f'/sys/devices/system/cpu/cpu{self.idx}/cpufreq/cpuinfo_cur_freq'
		output = sp.check_output(['adb', '-s', self.ip, 'shell', 'su -c', '\"cat', fname+'\"'])
		freq   = int(output.decode('utf-8').strip()) / 1e6
		# print(f'[c{self.idx}] current freq: {freq}')
		self.freq_data.append(freq)
	
	def get_freq_idx(self):
		'''
		Get the frequency level of current CPU instance
		'''
		fname  = f'/sys/devices/system/cpu/cpu{self.idx}/cpufreq/cpuinfo_cur_freq'
		output = sp.check_output(['adb', '-s', self.ip, 'shell', 'su -c', '\"cat', fname+'\"'])
		freq   = int(output.decode('utf-8').strip())
		return self.freq_list.index(freq)# TODO: / (len(self.freq_list) - 1)

	###############################
	# Static methods for all CPUs #
	###############################
	@staticmethod
	def print_available_freq(ip: str):
		'''
		Print all CPUs' available frequencies
		Pixel 4:
        	CPU[0,4): 300000, 403200, 499200, 576000, 672000, 768000, 844800, 940800, 1036800, 1113600, 1209600, 1305600, 1382400, 1478400, 1555200, 1632000, 1708800, 1785600
        	CPU[4,7): 710400, 825600, 940800, 1056000, 1171200, 1286400, 1401600, 1497600, 1612800, 1708800, 1804800, 1920000, 2016000, 2131200, 2227200, 2323200, 2419200
        	CPU[7]:   825600, 940800, 1056000, 1171200, 1286400, 1401600, 1497600, 1612800, 1708800, 1804800, 1920000, 2016000, 2131200, 2227200, 2323200, 2419200, 2534400, 2649600, 2745600, 2841600
		'''
		for i in range(CPU.n_cpu):
			fname  = f'/sys/devices/system/cpu/cpu{i}/cpufreq/scaling_available_frequencies'
			output = sp.check_output(['adb', '-s', ip, 'shell', 'su -c', '\"cat', fname+'\"'])
			output = output.decode('utf-8').strip()
			print(f'[CPU{i}] {output}')

	@staticmethod
	def set_scaling_governor(ip: str, gov: str):
		'''
		Set all CPUs' scaling governors to gov
			options: userspace, powersave, performance, schedutil
		'''
		# Check if the input gov is legal
		fname  = f'/sys/devices/system/cpu/cpu0/cpufreq/scaling_available_governors'
		output = sp.check_output(['adb', '-s', ip, 'shell', 'su -c', '\"cat', fname+'\"'])
		avail_govs = output.decode('utf-8').split()
		if gov not in avail_govs:
			raise ValueError(f'governor must be one of {avail_govs}')
		# Set all CPUs' governors
		for i in range(CPU.n_cpu):
			# Set the governor
			fname = f'/sys/devices/system/cpu/cpu{i}/cpufreq/scaling_governor'
			sp.check_output(['adb', '-s', ip, 'shell', 'su -c', '\"echo', gov+' >', fname+'\"'])
			# Print prompt
			output = sp.check_output(['adb', '-s', ip, 'shell', 'su -c', '\"cat', fname+'\"'])
		print(f'[CPU] set the governor to {output.decode("utf-8").strip()}')

	@staticmethod
	def print_all_freq(ip: str):
		'''
		Get all CPUs' current frequencies
		'''
		for i in range(CPU.n_cpu):
			fname  = f'/sys/devices/system/cpu/cpu{i}/cpufreq/cpuinfo_cur_freq'
			output = sp.check_output(['adb', '-s', ip, 'shell', 'su -c', '\"cat', fname+'\"'])
			print(f'[CPU{i}] {output.decode("utf-8").strip()} KHz')

	@staticmethod
	def get_overall_temp(ip: str):
		'''
		Get the temperature
		'''
		# TODO: thermal file
		# 输出所有的termal_zone的温度来比对
		fname  = '/sys/devices/virtual/thermal/thermal_zone10/temp'
		output = sp.check_output(['adb', '-s', ip, 'shell', 'su -c', '\"cat', fname+'\"'])
		temp   = int(output.decode('utf-8').strip()) / 1e3
		print(f'[CPU] temperature: {temp}')
		CPU.temp_data.append(temp)
		return temp